# -*- coding: utf-8 -*-
"""YOLO目标检测和跟踪引擎"""

import os
import cv2
import numpy as np
import logging
from ultralytics import YOLO
from collections import deque
import time

logger = logging.getLogger(__name__)

class DetectionEngine:
    """YOLO目标检测和跟踪引擎类"""
    
    def __init__(self):
        """初始化检测引擎"""
        self.model = None
        self.config = {}
        self.is_initialized = False
        self.tracker = None
        self.track_history = {}
        self.class_names = {}
        self.track_ids = 0
        self.last_track_time = time.time()
    
    def initialize(self, config):
        """初始化检测引擎
        
        Args:
            config: 检测引擎配置字典
            
        Returns:
            bool: 初始化是否成功
        """
        try:
            logger.info("初始化检测引擎...")
            self.config = config
            
            # 加载模型
            model_path = config.get('model_path', 'models/yolov8n.pt')
            if not os.path.exists(model_path):
                logger.error(f"模型文件不存在: {model_path}")
                return False
            
            # 加载YOLO模型
            self.model = YOLO(model_path)
            logger.info(f"YOLO模型加载成功: {model_path}")
            
            # 获取类别名称
            self.class_names = self.model.names
            logger.info(f"加载了 {len(self.class_names)} 个类别")
            
            # 初始化跟踪器
            self._initialize_tracker(config.get('tracking', {}))
            
            self.is_initialized = True
            logger.info("检测引擎初始化成功")
            return True
            
        except Exception as e:
            logger.error(f"初始化检测引擎失败: {str(e)}")
            return False
    
    def _initialize_tracker(self, tracking_config):
        """初始化目标跟踪器
        
        Args:
            tracking_config: 跟踪器配置
        """
        self.track_history = {}
        self.track_ids = 0
        self.tracking_config = tracking_config
        logger.info("目标跟踪器初始化完成")
    
    def detect(self, frame):
        """对图像帧进行目标检测
        
        Args:
            frame: 输入图像帧
            
        Returns:
            list: 检测结果列表
        """
        if not self.is_initialized:
            logger.error("检测引擎未初始化")
            return []
        
        try:
            # 获取配置参数
            confidence_threshold = self.config.get('confidence_threshold', 0.5)
            iou_threshold = self.config.get('iou_threshold', 0.45)
            device = self.config.get('device', 'cpu')
            classes = self.config.get('classes', None)
            
            # 执行检测
            results = self.model.predict(
                source=frame,
                conf=confidence_threshold,
                iou=iou_threshold,
                device=device,
                classes=classes,
                verbose=False
            )
            
            detections = []
            if results and len(results) > 0:
                result = results[0]
                
                # 解析检测结果
                for box in result.boxes:
                    x1, y1, x2, y2 = map(int, box.xyxy[0])
                    confidence = float(box.conf[0])
                    class_id = int(box.cls[0])
                    class_name = self.class_names.get(class_id, f'class_{class_id}')
                    
                    # 计算中心点
                    cx = (x1 + x2) // 2
                    cy = (y1 + y2) // 2
                    
                    detections.append({
                        'bbox': [x1, y1, x2, y2],
                        'confidence': confidence,
                        'class_id': class_id,
                        'class_name': class_name,
                        'center': (cx, cy),
                        'area': (x2 - x1) * (y2 - y1)
                    })
            
            return detections
            
        except Exception as e:
            logger.error(f"检测失败: {str(e)}")
            return []
    
    def track(self, detections, frame):
        """对检测结果进行目标跟踪
        
        Args:
            detections: 检测结果列表
            frame: 当前图像帧
            
        Returns:
            list: 跟踪结果列表
        """
        if not self.is_initialized:
            logger.error("检测引擎未初始化")
            return []
        
        try:
            tracked_objects = []
            current_time = time.time()
            
            # 使用简单的IoU跟踪算法
            # 1. 获取当前帧中的所有检测框
            current_boxes = [d['bbox'] for d in detections]
            current_centers = [d['center'] for d in detections]
            
            # 2. 更新已有跟踪器或创建新的跟踪器
            new_track_history = {}
            used_detection_indices = set()
            
            # 遍历历史跟踪对象
            for track_id, track_info in self.track_history.items():
                last_center = track_info['last_center']
                last_bbox = track_info['last_bbox']
                last_time = track_info['last_time']
                
                # 计算与当前检测的距离
                distances = []
                for i, (bbox, center) in enumerate(zip(current_boxes, current_centers)):
                    if i not in used_detection_indices:
                        # 使用中心点距离和IoU结合的方式
                        dist = np.sqrt((center[0] - last_center[0])**2 + (center[1] - last_center[1])** 2)
                        iou_val = self._calculate_iou(bbox, last_bbox)
                        
                        # 距离权重和IoU权重
                        distances.append((i, dist, iou_val))
                
                # 找到最佳匹配
                if distances:
                    # 根据距离和IoU排序，选择最佳匹配
                    distances.sort(key=lambda x: (x[1], -x[2]))
                    best_idx, best_dist, best_iou = distances[0]
                    
                    max_distance = 100  # 最大允许距离
                    min_iou = 0.1  # 最小允许IoU
                    
                    if best_dist < max_distance or best_iou > min_iou:
                        # 更新跟踪信息
                        detection = detections[best_idx]
                        track_info['last_bbox'] = detection['bbox']
                        track_info['last_center'] = detection['center']
                        track_info['last_time'] = current_time
                        track_info['class_id'] = detection['class_id']
                        track_info['class_name'] = detection['class_name']
                        track_info['confidence'] = detection['confidence']
                        
                        # 更新轨迹历史
                        track_info['trajectory'].append(detection['center'])
                        if len(track_info['trajectory']) > 100:  # 只保留最近100个点
                            track_info['trajectory'].pop(0)
                        
                        # 计算移动方向和速度
                        if len(track_info['trajectory']) > 1:
                            prev_center = track_info['trajectory'][-2]
                            current_center = track_info['trajectory'][-1]
                            track_info['direction'] = (
                                current_center[0] - prev_center[0],
                                current_center[1] - prev_center[1]
                            )
                            
                            time_diff = current_time - track_info['times'][-1]
                            if time_diff > 0:
                                track_info['speed'] = np.sqrt(
                                    (track_info['direction'][0] / time_diff)**2 +
                                    (track_info['direction'][1] / time_diff)** 2
                                )
                            track_info['times'].append(current_time)
                            if len(track_info['times']) > 100:
                                track_info['times'].pop(0)
                        
                        new_track_history[track_id] = track_info
                        used_detection_indices.add(best_idx)
                        
                        # 添加到跟踪结果
                        tracked_objects.append({
                            'track_id': track_id,
                            **detection,
                            'trajectory': track_info['trajectory'],
                            'direction': track_info.get('direction', (0, 0)),
                            'speed': track_info.get('speed', 0),
                            'tracking_time': current_time - track_info['start_time']
                        })
            
            # 为未匹配的检测创建新的跟踪器
            max_age = self.tracking_config.get('max_age', 30)
            for i, detection in enumerate(detections):
                if i not in used_detection_indices:
                    # 创建新的跟踪ID
                    self.track_ids += 1
                    track_id = self.track_ids
                    
                    # 初始化跟踪信息
                    track_info = {
                        'last_bbox': detection['bbox'],
                        'last_center': detection['center'],
                        'last_time': current_time,
                        'start_time': current_time,
                        'class_id': detection['class_id'],
                        'class_name': detection['class_name'],
                        'confidence': detection['confidence'],
                        'trajectory': [detection['center']],
                        'times': [current_time],
                        'direction': (0, 0),
                        'speed': 0
                    }
                    
                    new_track_history[track_id] = track_info
                    
                    # 添加到跟踪结果
                    tracked_objects.append({
                        'track_id': track_id,
                        **detection,
                        'trajectory': track_info['trajectory'],
                        'direction': track_info['direction'],
                        'speed': track_info['speed'],
                        'tracking_time': 0
                    })
            
            # 清理过期的跟踪器
            self.track_history = new_track_history
            
            # 定期清理长时间未更新的跟踪器
            if current_time - self.last_track_time > 1.0:  # 每秒清理一次
                self._cleanup_old_tracks(current_time, max_age)
                self.last_track_time = current_time
            
            return tracked_objects
            
        except Exception as e:
            logger.error(f"跟踪失败: {str(e)}")
            return []
    
    def _calculate_iou(self, box1, box2):
        """计算两个边界框的IoU
        
        Args:
            box1: [x1, y1, x2, y2]
            box2: [x1, y1, x2, y2]
            
        Returns:
            float: IoU值
        """
        x1 = max(box1[0], box2[0])
        y1 = max(box1[1], box2[1])
        x2 = min(box1[2], box2[2])
        y2 = min(box1[3], box2[3])
        
        # 计算交集面积
        intersection = max(0, x2 - x1) * max(0, y2 - y1)
        
        # 计算并集面积
        area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
        area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
        union = area1 + area2 - intersection
        
        if union == 0:
            return 0
        
        return intersection / union
    
    def _cleanup_old_tracks(self, current_time, max_age):
        """清理过期的跟踪器
        
        Args:
            current_time: 当前时间
            max_age: 最大允许的未更新时间（秒）
        """
        expired_ids = []
        for track_id, track_info in self.track_history.items():
            if current_time - track_info['last_time'] > max_age:
                expired_ids.append(track_id)
        
        for track_id in expired_ids:
            del self.track_history[track_id]
        
        if expired_ids:
            logger.debug(f"清理了 {len(expired_ids)} 个过期跟踪器")
    
    def detect_and_track(self, frame):
        """执行检测和跟踪
        
        Args:
            frame: 输入图像帧
            
        Returns:
            tuple: (detections, tracked_objects)
        """
        detections = self.detect(frame)
        tracked_objects = self.track(detections, frame)
        return detections, tracked_objects
    
    def get_tracking_stats(self):
        """获取跟踪统计信息
        
        Returns:
            dict: 跟踪统计信息
        """
        stats = {
            'active_tracks': len(self.track_history),
            'track_ids_count': self.track_ids,
            'track_history_size': sum(len(v['trajectory']) for v in self.track_history.values())
        }
        return stats
    
    def visualize(self, frame, detections=None, tracked_objects=None, show_detections=True, show_tracks=True):
        """可视化检测和跟踪结果
        
        Args:
            frame: 输入图像帧
            detections: 检测结果列表
            tracked_objects: 跟踪结果列表
            show_detections: 是否显示检测框
            show_tracks: 是否显示跟踪轨迹
            
        Returns:
            np.ndarray: 可视化后的图像
        """
        try:
            # 复制图像以避免修改原图
            visualized_frame = frame.copy()
            
            # 定义颜色映射
            colors = {
                0: (255, 0, 0),  # person - 蓝色
                1: (0, 255, 0),  # bicycle - 绿色
                2: (0, 0, 255),  # car - 红色
                3: (255, 255, 0),  # motorbike - 青色
                5: (255, 0, 255),  # bus - 紫色
                7: (0, 255, 255),  # truck - 黄色
                'default': (128, 128, 128)  # 默认 - 灰色
            }
            
            # 显示跟踪结果
            if show_tracks and tracked_objects:
                for track_obj in tracked_objects:
                    track_id = track_obj['track_id']
                    bbox = track_obj['bbox']
                    class_id = track_obj['class_id']
                    class_name = track_obj['class_name']
                    confidence = track_obj['confidence']
                    trajectory = track_obj.get('trajectory', [])
                    
                    # 获取颜色
                    color = colors.get(class_id, colors['default'])
                    
                    # 绘制边界框
                    x1, y1, x2, y2 = bbox
                    cv2.rectangle(visualized_frame, (x1, y1), (x2, y2), color, 2)
                    
                    # 绘制标签
                    label = f"ID:{track_id} {class_name} {confidence:.2f}"
                    cv2.putText(visualized_frame, label, (x1, y1 - 10),
                               cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
                    
                    # 绘制轨迹
                    if len(trajectory) > 1:
                        for i in range(1, len(trajectory)):
                            cv2.line(visualized_frame, trajectory[i-1], trajectory[i],
                                    color, 2, cv2.LINE_AA)
                    
                    # 绘制中心点
                    center = track_obj['center']
                    cv2.circle(visualized_frame, center, 4, color, -1)
            
            # 显示检测结果（如果没有跟踪结果）
            elif show_detections and detections:
                for detection in detections:
                    bbox = detection['bbox']
                    class_id = detection['class_id']
                    class_name = detection['class_name']
                    confidence = detection['confidence']
                    
                    # 获取颜色
                    color = colors.get(class_id, colors['default'])
                    
                    # 绘制边界框
                    x1, y1, x2, y2 = bbox
                    cv2.rectangle(visualized_frame, (x1, y1), (x2, y2), color, 2)
                    
                    # 绘制标签
                    label = f"{class_name} {confidence:.2f}"
                    cv2.putText(visualized_frame, label, (x1, y1 - 10),
                               cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
            
            return visualized_frame
            
        except Exception as e:
            logger.error(f"可视化失败: {str(e)}")
            return frame
    
    def shutdown(self):
        """关闭检测引擎"""
        try:
            logger.info("关闭检测引擎...")
            self.model = None
            self.track_history.clear()
            self.is_initialized = False
            logger.info("检测引擎已关闭")
            
        except Exception as e:
            logger.error(f"关闭检测引擎时发生错误: {str(e)}")
    
    def is_ready(self):
        """检查引擎是否准备就绪
        
        Returns:
            bool: 是否就绪
        """
        return self.is_initialized and self.model is not None

# 创建全局检测引擎实例
global_detection_engine = None

def get_detection_engine():
    """获取全局检测引擎实例
    
    Returns:
        DetectionEngine: 检测引擎实例
    """
    global global_detection_engine
    if global_detection_engine is None:
        global_detection_engine = DetectionEngine()
    return global_detection_engine